from sentence_transformers import SentenceTransformer
from configs import TEXT2VEC_LARGE_MODEL_PATH, device
from faiss import normalize_L2
from tqdm import trange
import numpy as np
import torch
import logging
from torch import Tensor
logger = logging.getLogger(__name__)
# TEXT2VEC_LARGE_MODEL_PATH = "../embedding_models/text2vec_large"
# device = torch.device("cuda:3")
show_progress_bar = (logger.getEffectiveLevel()==logging.INFO or logger.getEffectiveLevel()==logging.DEBUG)
model = SentenceTransformer(TEXT2VEC_LARGE_MODEL_PATH)
model.eval()
model = model.to(device)


def batch_to_device(batch):
    for key in batch:
        if isinstance(batch[key], Tensor):
            batch[key] = batch[key].to(device)
    return batch


def text2vec_encode(text_list, is_normalize=True):
    batch_size = 32
    if not isinstance(text_list, list):
        text_list = [text_list]
        flag = 0
    else:
        flag = 1
    length_sorted_idx = np.argsort([-len(sen) for sen in text_list])
    sentences_sorted = [text_list[idx] for idx in length_sorted_idx]
    all_embeddings = []
    with torch.no_grad():
        for start_index in trange(0, len(text_list), batch_size, desc="Batches", disable=not show_progress_bar):
            sentences_batch = sentences_sorted[start_index:start_index+batch_size]
            encoded_input = model.tokenize(sentences_batch)
            encoded_input = batch_to_device(encoded_input)
            fea = model(encoded_input)
            embeddings = fea["sentence_embedding"]
            embeddings = embeddings.detach()
            if is_normalize:
                embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
            embeddings = embeddings.cpu()
            all_embeddings.extend(embeddings)
    all_embeddings = [all_embeddings[idx] for idx in np.argsort(length_sorted_idx)]
    all_embeddings = np.asarray([emb.numpy() for emb in all_embeddings])
    if flag == 0:
        return all_embeddings[0]
    else:
        return all_embeddings
    return all_embeddings



embed = text2vec_encode("被保险人被宣告死亡时如何处理?")
print(embed.shape)